from catboost import CatBoostRegressor
from sklearn.model_selection import GridSearchCV
from dataPretreat import data_pretreat


if __name__ == '__main__':
    # 加载数据

    iszt_train = data_pretreat(path='../resource/KDDTrain+.csv')
    iszt_test = data_pretreat(path='../resource/KDDTest+.csv')
    # 划分训练集和测试集

    X_train = iszt_train[
        ['duration', 'protocol_type', 'service', 'src_bytes', 'dst_bytes', 'wrong_fragment', 'serror_rate',
         'dst_host_srv_count', 'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate',
         'dst_host_rerror_rate']]
    y_train = iszt_train['label']
    X_test = iszt_test[
        ['duration', 'protocol_type', 'service', 'src_bytes', 'dst_bytes', 'wrong_fragment', 'serror_rate',
         'dst_host_srv_count', 'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate',
         'dst_host_rerror_rate']]
    y_test = iszt_test['label']

    # 转换为Dataset数据格式

    other_params = {
        'iterations': 1000,
        'learning_rate': 0.03,
        'l2_leaf_reg': 3,
        'bagging_temperature': 1,
        'random_strength': 1,
        'depth': 6,
        'rsm': 1,
        'one_hot_max_size': 2,
        'leaf_estimation_method': 'Gradient',
        'fold_len_multiplier': 2,
        'border_count': 128,
    }
    model_cb = CatBoostRegressor(**other_params)
    optimized_cb = GridSearchCV(estimator=model_cb, param_grid=cv_params, scoring='r2', cv=5, verbose=1, n_jobs=2)
    optimized_cb.fit(X_train, y_train, cat_features=category_features)
    print('参数的最佳取值：{0}'.format(optimized_cb.best_params_))
    print('最佳模型得分:{0}'.format(optimized_cb.best_score_))
    print(optimized_cb.cv_results_['mean_test_score'])
    print(optimized_cb.cv_results_['params'])
